#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/spi/spi.h>
#include <linux/fb.h>
#include <linux/delay.h>
#include <linux/gpio/consumer.h>
#include <linux/of.h>
#include <linux/of_device.h>
#include <linux/types.h>
#include <linux/string.h>
#include <linux/dma-mapping.h>
#include <linux/vmalloc.h>
#include <linux/slab.h>
#include <linux/mutex.h>

#define DRV_NAME "st7789fb"

// ST7789分辨率
#define ST7789_WIDTH  240
#define ST7789_HEIGHT 240

#define ST7789_FPS     30
// ST7789命令
#define ST7789_NOP     0x00
#define ST7789_SWRESET 0x01
#define ST7789_SLPOUT  0x11
#define ST7789_INVOFF  0x20
#define ST7789_INVON   0x21
#define ST7789_DISPON  0x29
#define ST7789_CASET   0x2A
#define ST7789_RASET   0x2B
#define ST7789_RAMWR   0x2C
#define ST7789_MADCTL  0x36
#define ST7789_COLMOD  0x3A

// 颜色格式配置 (16位RGB565)
#define ST7789_COLOR_MODE_16BIT 0x55

struct st7789fb_dev {
    struct spi_device *spi;
    struct fb_info *info;
    struct gpio_desc *dc;           // 数据/命令引脚
    struct gpio_desc *reset;        // 复位引脚
    struct mutex lock;              // 保护 fb 数据访问
    struct task_struct *thread;     // 刷新线程
    dma_addr_t dma_addr;            // DMA 地址（可选）
    u16 *buffer;                    // 显示缓冲区
    u32 buffer_size;                // 缓冲区大小
};

// 向ST7789发送命令
static void st7789_send_cmd(struct st7789fb_dev *dev, u8 cmd)
{

    int ret;
    gpiod_set_value(dev->dc, 0); // 命令模式
    ret = spi_write(dev->spi, &cmd, 1);
    if (ret < 0)
        dev_err(&dev->spi->dev, "Failed to send command 0x%02x: %d\n", cmd, ret);
}

// 向ST7789发送数据
static void st7789_send_data(struct st7789fb_dev *dev, const u8 *data, size_t len)
{
    int ret;
    gpiod_set_value(dev->dc, 1); // 数据模式
    ret = spi_write(dev->spi, data, len);
    if (ret < 0)
        dev_err(&dev->spi->dev, "Failed to send data: %d\n", ret);
}

static void st7789_send_data_byte(struct st7789fb_dev *dev, const u8 data)
{
    st7789_send_data(dev, &data, 1);
}

// 带重试机制的SPI写入
static int st7789_spi_write_retry(struct spi_device *spi, const void *buf, size_t len)
{
    int ret;
    int retry = 3;
    
    while (retry--) {
        ret = spi_write(spi, buf, len);
        if (ret >= 0)
            return ret;
        msleep(10);
    }
    return ret;
}



// 初始化ST7789
static int st7789_init(struct st7789fb_dev *dev)
{
    // 复位显示屏
    if (dev->reset) {
        gpiod_set_value(dev->reset, 0);
        msleep(100);
        gpiod_set_value(dev->reset, 1);
        msleep(100);
    }

    // 软件复位
    st7789_send_cmd(dev, ST7789_SWRESET);
    msleep(150);

    // 初始化序列
    st7789_send_cmd(dev, ST7789_SLPOUT); // Sleep out
    msleep(120);
    
    st7789_send_cmd(dev, ST7789_MADCTL);
    st7789_send_data_byte(dev, 0x00);  // 根据硬件调整
    
    st7789_send_cmd(dev, ST7789_COLMOD);
    st7789_send_data_byte(dev, ST7789_COLOR_MODE_16BIT);
    
    // 电源及显示配置
    st7789_send_cmd(dev, 0xB2);
    st7789_send_data_byte(dev, 0x1F);
    st7789_send_data_byte(dev, 0x1F);
    st7789_send_data_byte(dev, 0x00);
    st7789_send_data_byte(dev, 0x33);
    st7789_send_data_byte(dev, 0x33);
    
    st7789_send_cmd(dev, 0xB7);
    st7789_send_data_byte(dev, 0x35);
    
    st7789_send_cmd(dev, 0xBB);
    st7789_send_data_byte(dev, 0x2B);
    
    st7789_send_cmd(dev, 0xC0);
    st7789_send_data_byte(dev, 0x2C);
    
    st7789_send_cmd(dev, 0xC2);
    st7789_send_data_byte(dev, 0x01);
    
    st7789_send_cmd(dev, 0xC3);
    st7789_send_data_byte(dev, 0x0F);
    
    st7789_send_cmd(dev, 0xC4);
    st7789_send_data_byte(dev, 0x20);
    
    st7789_send_cmd(dev, 0xC6);
    st7789_send_data_byte(dev, 0x13);
    
    st7789_send_cmd(dev, 0xD0);
    st7789_send_data_byte(dev, 0xA4);
    st7789_send_data_byte(dev, 0xA1);
    
    st7789_send_cmd(dev, 0xD6);
    st7789_send_data_byte(dev, 0xA1);
    
    // Gamma校正
    st7789_send_cmd(dev, 0xE0);
    st7789_send_data_byte(dev, 0xF0);
    st7789_send_data_byte(dev, 0x04);
    st7789_send_data_byte(dev, 0x07);
    st7789_send_data_byte(dev, 0x04);
    st7789_send_data_byte(dev, 0x04);
    st7789_send_data_byte(dev, 0x04);
    st7789_send_data_byte(dev, 0x25);
    st7789_send_data_byte(dev, 0x33);
    st7789_send_data_byte(dev, 0x3C);
    st7789_send_data_byte(dev, 0x36);
    st7789_send_data_byte(dev, 0x14);
    st7789_send_data_byte(dev, 0x12);
    st7789_send_data_byte(dev, 0x29);
    st7789_send_data_byte(dev, 0x30);
    
    st7789_send_cmd(dev, 0xE1);
    st7789_send_data_byte(dev, 0xF0);
    st7789_send_data_byte(dev, 0x02);
    st7789_send_data_byte(dev, 0x04);
    st7789_send_data_byte(dev, 0x05);
    st7789_send_data_byte(dev, 0x05);
    st7789_send_data_byte(dev, 0x21);
    st7789_send_data_byte(dev, 0x25);
    st7789_send_data_byte(dev, 0x32);
    st7789_send_data_byte(dev, 0x3B);
    st7789_send_data_byte(dev, 0x38);
    st7789_send_data_byte(dev, 0x12);
    st7789_send_data_byte(dev, 0x14);
    st7789_send_data_byte(dev, 0x27);
    st7789_send_data_byte(dev, 0x31);
    
    st7789_send_cmd(dev, 0xE4);
    st7789_send_data_byte(dev, 0x1D);
    st7789_send_data_byte(dev, 0x00);
    st7789_send_data_byte(dev, 0x00);
    
    st7789_send_cmd(dev, ST7789_INVON);  // 打开反显（根据硬件调整）
    st7789_send_cmd(dev, ST7789_DISPON); // 显示开启
    
    return 0;
}

// 设置显示窗口
static void st7789_set_window(struct st7789fb_dev *dev, u16 x0, u16 y0, u16 x1, u16 y1)
{
    u8 data[4];

    // 设置列地址
    st7789_send_cmd(dev, ST7789_CASET);
    data[0] = (x0 >> 8) & 0xFF;
    data[1] = x0 & 0xFF;
    data[2] = (x1 >> 8) & 0xFF;
    data[3] = x1 & 0xFF;
    st7789_send_data(dev, data, 4);

    // 设置行地址
    st7789_send_cmd(dev, ST7789_RASET);
    data[0] = (y0 >> 8) & 0xFF;
    data[1] = y0 & 0xFF;
    data[2] = (y1 >> 8) & 0xFF;
    data[3] = y1 & 0xFF;
    st7789_send_data(dev, data, 4);

    // 准备写入GRAM
    st7789_send_cmd(dev, ST7789_RAMWR);
}

// 高效清屏函数（批量传输）
static void st7789_clear_screen(struct st7789fb_dev *dev, u16 color)
{
    const u32 total_pixels = ST7789_WIDTH * ST7789_HEIGHT;
    const u32 chunk_size = 4096 / sizeof(u16); // 4KB缓冲区
    u16 *temp_buf;
    u32 remaining = total_pixels;

    temp_buf = kmalloc(chunk_size * sizeof(u16), GFP_KERNEL);
    if (!temp_buf) {
        dev_err(&dev->spi->dev, "Failed to alloc clear buffer\n");
        return;
    }

    // 填充临时缓冲区
    memset(temp_buf, color, chunk_size * sizeof(u16));

    // 设置全屏窗口
    st7789_set_window(dev, 0, 0, ST7789_WIDTH - 1, ST7789_HEIGHT - 1);
    gpiod_set_value(dev->dc, 1);

    // 批量传输
    while (remaining > 0) {
        u32 send = min(remaining, chunk_size);
        st7789_spi_write_retry(dev->spi, temp_buf, send * sizeof(u16));
        remaining -= send;
    }

    kfree(temp_buf);
    // 更新本地缓冲区
    memset(dev->buffer, color, total_pixels * sizeof(u16));
}


// 验证显示参数
static int st7789fb_check_var(struct fb_var_screeninfo *var, struct fb_info *info)
{
    // 支持旋转角度检查
    if (var->rotate % 90 != 0 || var->rotate < 0 || var->rotate > 270)
        return -EINVAL;

    // 分辨率检查
    if (var->xres != ST7789_WIDTH || var->yres != ST7789_HEIGHT ||
        var->xres_virtual != var->xres || var->yres_virtual != var->yres)
        return -EINVAL;

    // 像素格式检查 (RGB565)
    if (var->bits_per_pixel != 16 ||
        var->red.length != 5 || var->red.offset != 11 ||
        var->green.length != 6 || var->green.offset != 5 ||
        var->blue.length != 5 || var->blue.offset != 0)
        return -EINVAL;

    return 0;
}

// 应用显示参数设置
static int st7789fb_set_par(struct fb_info *info)
{
    struct st7789fb_dev *dev = info->par;
    struct fb_var_screeninfo *var = &info->var;
    u8 madctl = 0x00;

    // 根据旋转角度设置MADCTL
    switch (var->rotate) {
        case 0:
            madctl = 0x00;  // 正常方向
            break;
        case 90:
            madctl = 0x60;  // 顺时针90度
            break;
        case 180:
            madctl = 0xC0;  // 180度
            break;
        case 270:
            madctl = 0xA0;  // 顺时针270度
            break;
        default:
            return -EINVAL;
    }

    st7789_send_cmd(dev, ST7789_MADCTL);
    st7789_send_data(dev, &madctl, 1);
    msleep(10);

    return 0;
}

// Framebuffer操作函数集
static struct fb_ops st7789fb_ops = {
    .owner = THIS_MODULE,
    .fb_fillrect = cfb_fillrect,
	.fb_imageblit = cfb_imageblit,
	.fb_copyarea = cfb_copyarea,
    .fb_check_var = st7789fb_check_var,
    .fb_set_par = st7789fb_set_par,
};

// 线程刷新屏幕
static int st7789_refresh_thread(void *data)
{
    
    struct st7789fb_dev *dev = data;

    int ret;
    unsigned int refresh_interval = (1000 / ST7789_FPS);  // 刷新间隔（ms）
   
    dev_info(&dev->spi->dev, "st7789_refresh_thread started (DMA mode, using pre-allocated dma_addr)\n");

    while (!kthread_should_stop()) {
        mutex_lock(&dev->lock);

        // 2. 设置全屏显示区域
        st7789_set_window(dev, 0, 0, ST7789_WIDTH-1, ST7789_HEIGHT-1);

        struct spi_transfer t = {
            .tx_buf = dev->buffer,      // 虚拟地址（CPU可见）
            .tx_dma = dev->dma_addr,    // DMA地址（设备可见）
            .len    = dev->buffer_size,
        };
        struct spi_message m;
        spi_message_init(&m);
        spi_message_add_tail(&t, &m);

        // 4. 执行SPI DMA传输
        gpiod_set_value(dev->dc, 1);
        ret = spi_sync(dev->spi, &m);
        if (ret) {
            dev_err(&dev->spi->dev, "SPI DMA transfer failed: %d\n", ret);
        }

        mutex_unlock(&dev->lock);
        msleep(refresh_interval);
    }

    return 0;
}


// 探针函数
static int st7789fb_probe(struct spi_device *spi)
{
    struct st7789fb_dev *dev;
    struct fb_info *info;
    int ret;

    // 分配设备结构体
    dev = devm_kzalloc(&spi->dev, sizeof(*dev), GFP_KERNEL);
    if (!dev)
        return -ENOMEM;

    mutex_init(&dev->lock);

    dev->spi = spi;
    spi_set_drvdata(spi, dev);

    // 获取GPIO引脚
    dev->dc = devm_gpiod_get(&spi->dev, "dc", GPIOD_OUT_LOW);
    if (IS_ERR(dev->dc)) {
        dev_err(&spi->dev, "Failed to get DC GPIO: %ld\n", PTR_ERR(dev->dc));
        return PTR_ERR(dev->dc);
    }

    dev->reset = devm_gpiod_get_optional(&spi->dev, "reset", GPIOD_OUT_HIGH);
    if (IS_ERR(dev->reset)) {
        dev_err(&spi->dev, "Failed to get RESET GPIO: %ld\n", PTR_ERR(dev->reset));
        return PTR_ERR(dev->reset);
    }

    /** 设置DMA掩码，确保支持32位地址 !!! 若不设置，dma_alloc_coherent分配内存必失败*/
    if (!spi->dev.coherent_dma_mask) {
        ret = dma_coerce_mask_and_coherent(&spi->dev, DMA_BIT_MASK(32));
        if (ret) {
            dev_err(&spi->dev, "Failed to set DMA mask: %d\n", ret);
            return ret;
        }
    }

    /**  
     * 两种选择配置SPI 参数
     * 1.代码直接配置，如下所示，会覆盖设备树的配置，只生效代码的配置
     * 2.推荐 修改设备树，在设备树里面配置SPI参数，注释掉代码配置
     *    设备树配置方式见 st7789fb.dts
     *    spi-max-frequency = <40000000>;
    */

    spi->mode = SPI_MODE_3;         // CPOL=1, CPHA=1
    spi->bits_per_word = 8;         // 8bit
    spi->max_speed_hz = 40000000;   // 40MHz

    ret = spi_setup(spi); 
    if (ret < 0) {
        dev_err(&spi->dev, "Failed to setup SPI: %d\n", ret);
        return ret;
    }

    // 初始化显示屏
    ret = st7789_init(dev);
    if (ret < 0) {
        dev_err(&spi->dev, "Failed to initialize ST7789: %d\n", ret);
        return ret;
    }

    // 分配Framebuffer信息结构体
    info = framebuffer_alloc(0, &spi->dev);
    if (!info)
        return -ENOMEM;

    dev->info = info;
    info->par = dev;
    info->fbops = &st7789fb_ops;

    // 配置显示参数
    info->var.xres = ST7789_WIDTH;
    info->var.yres = ST7789_HEIGHT;
    info->var.xres_virtual = ST7789_WIDTH;
    info->var.yres_virtual = ST7789_HEIGHT;
    info->var.bits_per_pixel = 16;
    info->var.rotate = 0;  // 默认不旋转

    // 设置像素格式 (RGB565)
    info->var.red.offset = 11;
    info->var.red.length = 5;
    info->var.green.offset = 5;
    info->var.green.length = 6;
    info->var.blue.offset = 0;
    info->var.blue.length = 5;

    info->fix.smem_len = ST7789_WIDTH * ST7789_HEIGHT * 2;
    info->fix.type = FB_TYPE_PACKED_PIXELS;
    info->fix.visual = FB_VISUAL_TRUECOLOR;
    info->fix.line_length = ST7789_WIDTH * 2;

    // 分配显示缓冲区（优先DMA，失败则用vmalloc）
    info->screen_base = dma_alloc_coherent(&spi->dev, info->fix.smem_len,
                                       (dma_addr_t*)&info->fix.smem_start, GFP_KERNEL);

    if (!info->screen_base) {
        dev_info(&spi->dev, "DMA allocation failed, trying vmalloc\n");
        ret = -ENOMEM;
        goto err_fb_release;
    }

    dev->dma_addr = info->fix.smem_start;
    dev->buffer = (u16 *)info->screen_base;
    dev->buffer_size = info->fix.smem_len;

    // 注册Framebuffer设备
    ret = register_framebuffer(info);
    if (ret < 0) {
        dev_err(&spi->dev, "Failed to register framebuffer: %d\n", ret);
        goto err_buffer_free;
    }

    // 创建刷新线程
    dev->thread = kthread_run(st7789_refresh_thread, dev, "st7789-refresh");
    if (IS_ERR(dev->thread)) {
        ret = PTR_ERR(dev->thread);
        unregister_framebuffer(dev->info);
        framebuffer_release(dev->info);
        dma_free_coherent(&spi->dev, info->fix.smem_len,
                            info->screen_base, info->fix.smem_start);

        dev_err(&spi->dev, "Failed to create refresh thread: %d\n", ret);
        return ret;
    }

    st7789_clear_screen(dev, 0xFFFF);  // 清屏为白色，测试spi屏幕驱动

    dev_info(&spi->dev, "ST7789 Framebuffer driver loaded, %dx%d@16bpp\n",
             ST7789_WIDTH, ST7789_HEIGHT);

    return 0;

err_buffer_free:
    dma_free_coherent(&spi->dev, info->fix.smem_len,
                        info->screen_base, info->fix.smem_start);
err_fb_release:
    framebuffer_release(info);
    return ret;
}

// 移除设备时调用
static void st7789fb_remove(struct spi_device *spi)
{
    struct st7789fb_dev *dev = spi_get_drvdata(spi);
    // 停止刷新线程
    kthread_stop(dev->thread);

    unregister_framebuffer(dev->info);
    if (dev->info->screen_base) {
        dma_free_coherent(&spi->dev, dev->info->fix.smem_len,
                        dev->info->screen_base, dev->info->fix.smem_start);

    }
    framebuffer_release(dev->info);
}

// 设备ID匹配
static const struct spi_device_id st7789fb_id[] = {
    { "st7789", 0 },
    { /* Sentinel */ }
};

// 设备树匹配表
static const struct of_device_id st7789fb_of_match[] = {
    { .compatible = "sitronix,st7789" },
    { /* Sentinel */ }
};
MODULE_DEVICE_TABLE(of, st7789fb_of_match);

// SPI驱动结构体
static struct spi_driver st7789fb_driver = {
    .driver = {
        .name = DRV_NAME,
        .of_match_table = st7789fb_of_match,
    },
    .probe = st7789fb_probe,
    .remove = st7789fb_remove,
    .id_table = st7789fb_id,
};
module_spi_driver(st7789fb_driver);

MODULE_AUTHOR("klein <dfy1999@foxmail.com>");
MODULE_DESCRIPTION("ST7789 SPI LCD Framebuffer Driver");
MODULE_LICENSE("GPL");
MODULE_ALIAS("spi:" DRV_NAME);